package top.babyyan.base.config;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.util.concurrent.ThreadPoolExecutor;

/**
 * @Author : MYX
 * @create 2024/9/8
 */
@Configuration
@EnableAsync
public class AsyncConfig {


    /**
     * 异步日志线程池
     *
     * @return
     */
    @Bean(name = "logAsyncServiceExecutor")
    public ThreadPoolTaskExecutor getAsyncThreadPoolTaskExecutor() {
        // 获取当前服务器的 CPU 核心数
        int cpuCores = Runtime.getRuntime().availableProcessors();

        // 动态计算核心线程数，通常核心线程数为 CPU 核心数的 2 倍
        // 核心线程数是线程池在正常情况下始终维持的最小线程数
        int corePoolSize = cpuCores * 2;

        // 动态计算最大线程数，通常设置为核心线程数的 2 倍或更多
        // 最大线程数是线程池能够容纳的最大线程数，当队列已满时，线程池会尝试扩展到这个数量
        int maxPoolSize = cpuCores * 4;

        // 设置任务队列的容量，当核心线程都在工作时，新的任务将进入这个队列等待执行
        // 队列的大小应该根据系统负载和业务需求来决定
        int queueCapacity = 500;

        // 设置空闲线程的存活时间（单位：秒）
        // 当线程池中线程的数量超过核心线程数，并且这些多余线程空闲超过指定时间后，线程将被回收
        int keepAliveSeconds = 60;

        // 创建 ThreadPoolTaskExecutor 实例
        ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();

        // 配置核心线程数
        taskExecutor.setCorePoolSize(corePoolSize);

        // 配置最大线程数
        taskExecutor.setMaxPoolSize(maxPoolSize);

        // 配置队列容量
        taskExecutor.setQueueCapacity(queueCapacity);

        // 配置线程空闲存活时间
        taskExecutor.setKeepAliveSeconds(keepAliveSeconds);

        // 设置线程名称的前缀，方便问题排查时识别线程池中的线程
        taskExecutor.setThreadNamePrefix("async-log-service-");

        // 设置拒绝策略，当任务数超过最大线程数且队列已满时，采取的处理方式
        // CallerRunsPolicy：由调用线程（提交任务的线程）处理该任务，这样可以防止任务丢失
        taskExecutor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());

        // 设置线程池关闭时是否等待所有任务完成后再关闭
        taskExecutor.setWaitForTasksToCompleteOnShutdown(true);

        // 设置等待任务完成的最长时间（单位：秒），防止线程池长时间无法关闭
        taskExecutor.setAwaitTerminationSeconds(keepAliveSeconds);

        // 初始化线程池
        taskExecutor.initialize();

        // 返回配置好的线程池
        return taskExecutor;
    }

    /**
     * 应用公共线程池
     *
     * @return
     */
    @Bean(name = "threadPoolTaskExecutor")
    public ThreadPoolTaskExecutor threadPoolTaskExecutor() {
        // 获取当前服务器的 CPU 核心数
        int cpuCores = Runtime.getRuntime().availableProcessors();

        // 动态计算核心线程数和最大线程数
        // 这里核心线程数为 CPU 核心数的 4 倍，适合更大并发需求
        int corePoolSize = cpuCores * 4;
        // 最大线程数为核心线程数的 2.5 倍
        int maxPoolSize = cpuCores * 10;
        int queueCapacity = 1000;
        int keepAliveSeconds = 60;

        // 创建 ThreadPoolTaskExecutor 实例
        ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();

        // 配置核心线程数
        taskExecutor.setCorePoolSize(corePoolSize);

        // 配置最大线程数
        taskExecutor.setMaxPoolSize(maxPoolSize);

        // 配置队列容量
        taskExecutor.setQueueCapacity(queueCapacity);

        // 配置线程空闲时间
        taskExecutor.setKeepAliveSeconds(keepAliveSeconds);

        // 设置线程名称的前缀，方便识别线程
        taskExecutor.setThreadNamePrefix("method-exec-thread-pool-");

        // 设置拒绝策略，调用者线程运行任务，防止任务丢失
        taskExecutor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());

        // 配置关闭时等待所有任务完成再关闭线程池
        taskExecutor.setWaitForTasksToCompleteOnShutdown(true);

        // 以下设置是为了解决线程池环境下 RequestAttributes 获取错误的问题
        // parallelStream 默认使用公共线程池，可能导致 ThreadLocal 信息丢失
        taskExecutor.setTaskDecorator(runnable -> {
            // 在主线程中获取当前请求的属性
            RequestAttributes requestAttributes = RequestContextHolder.currentRequestAttributes();
            return () -> {
                try {
                    // 在执行任务的线程中设置请求属性
                    RequestContextHolder.setRequestAttributes(requestAttributes);
                    // 执行实际的任务
                    runnable.run();
                } finally {
                    // 任务执行完成后，重置请求属性，防止内存泄漏
                    RequestContextHolder.resetRequestAttributes();
                }
            };
        });

        // 初始化线程池
        taskExecutor.initialize();
        return taskExecutor;
    }
}

